import json
import sys
import logging
from datetime import datetime

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.models import resnet18
from tqdm import tqdm
from torch.utils.data import DataLoader
import torchvision
import torchvision.transforms as transforms

from _dynamic_dataset.AccCombinationRecord import AccCombinationRecord
from _dynamic_dataset.EvalDataset import EvalDataset
from _dynamic_dataset.EvalSubset import EvalSubset
from _dynamic_dataset.EvalType import EvalType
from algorithm.MR_DD import Algorithm_MR_DD_ACC_COMPARE
from rootPath import project_path
from server_client.Client import Client
from server_client.Server import Server
from utils.datasetUtils import get_SDSS_clients
from utils.modelUtils import getResNet18, getAvailableDevice, getResNet50
from utils.randomUtils import all_combinations

# 超参数
batch_size = 128
learning_rate = 0.0001
num_epochs = 250
threshold = 1.00
# 基于梯度的夏普利值计算
with open(project_path + "/conf/client.json", 'r') as f:
    clientConf = json.load(f)
with open(project_path + "/conf/server.json", 'r') as f:
    serverConf = json.load(f)
# 配置日志记录器
# # 配置日志记录器
# 获取当前日期和时间
now = datetime.now()
dirPath = project_path + "/experiment/mr_dd/logs/sdss2" + f"/{now.year}_{now.month:02d}_{now.day:02d}_{now.hour:02d}_{now.minute:02d}_{now.second:02d}"

# 数据预处理
train_transform = transforms.Compose(
    [
     transforms.RandomCrop(32, padding=4),
     transforms.RandomHorizontalFlip(),
     transforms.ToTensor(),
     transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])

# 测试集数据预处理
test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

# CIFAR-10数据集
trainset = torchvision.datasets.CIFAR10(root=project_path+'/data', train=True, download=True, transform=train_transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=2)

# 服务器的动态剪枝验证集
val_set = torchvision.datasets.CIFAR10(root=project_path + '/data', train=False, download=True,
                                       transform=test_transform)
val_set2 = torchvision.datasets.CIFAR10(root=project_path + '/data', train=False, download=True,
                                       transform=test_transform)
real_eval_loader = DataLoader(val_set2, batch_size=batch_size, shuffle=False, num_workers=2)
evalDataset = EvalDataset(val_set.data, val_set.targets, 30, 1.00, test_transform)
evalSubDataset = EvalSubset(evalDataset)
eval_loader = DataLoader(evalSubDataset, batch_size=batch_size, shuffle=False, num_workers=2)
# 测试用例总数
test_example_num = len(val_set.targets)
# 待训练模型
model = getResNet50()
device = getAvailableDevice()
clientNum = 5
clientIds = []
for i in range(clientNum):
    clientIds.append(i)
combinations = all_combinations(clientIds)
# 服务器和客户端
server = Server(serverConf, eval_loader, device, model, real_eval_loader)
# 创建相同分布相同大小的客户端
# client = Client(clientConf, model, device,trainloader, 0)
# clients = [client]
clients = get_SDSS_clients(clientConf, model, device, train_transform)
# 训练模型
acc_record = AccCombinationRecord()
fit_acc_record = AccCombinationRecord()
ignore_record = AccCombinationRecord()
sample_ratio_list = [0, 0.1, 0.1, 0.1, 0.1, 0.1]
soft_beta = 1.0
Algorithm_MR_DD_ACC_COMPARE(num_epochs, clients, server, device, soft_beta, sample_ratio_list, combinations, 10000, dirPath, acc_record,
                fit_acc_record, ignore_record, EvalType.RANDOM_CHOOSE_IGNORE_SET)
